library(Seurat)
library(Matrix)

args=commandArgs(T)

workDir=args[1]
CDNAlist=args[2] #third col is citeseq's num.0 ~ no citeseq,x,how many citeseq barcode
sample=args[3]
flag=args[4] #1,raw matrix;2,normal filter matrix;3,normal + add special filter matrix 4,only output normal filter matrix's barcode. 5. only use special barcode
ENSGflag=args[5] #T,the row names is ENSG
OutBarFlag=args[6] # T,output this gcmat's barcode List
SpeBarFile=args[7] #if add special barcode.no fill in F

#from cellrange's analysis data get GC Matrix
#cite seq is ok! read from CDNA.list.txt
#barcode prefix name come from CDNA list first col

#finish in 2020.10.14
#CITESEQ mul samples mul barcode is ok!

#output:
#if citeseq:P28.GCMat.RData,P28.CITESEQBarExpMat.RData,P28.CITEBarcodeExpTable.RData
#if not citeseq:P28.GCMat.RData

#eg:Rscript getFinalGCMat1.R /home/xingjun/zdata1/AnaTemp1/NormalFilter /data2/users/dongxingjun/Human_SC_Ana/P1/P1.CDNA.list P1 4 2 NULL

#------------------

#function
readCDNAList <- function(LTable){
  a <- read.table(LTable,sep = "\t")
  sampleList  <- as.vector(a$V1)
  CDNAdir <- as.vector(a$V2)
  CSflag <- as.vector(a$V3)

  names(CDNAdir) <- sampleList
  names(CSflag) <- sampleList

  ff <- list(sampleList,CDNAdir,CSflag)
  return(ff)
}

mergeTwoMat <- function(AMat,BMat)
{
  RMatA <- rownames(AMat)
  CMatA <- colnames(AMat)
  
  RMatB <- rownames(BMat)
  CMatB <- colnames(BMat)
  
  ABCommR <- intersect(RMatA,RMatB)
  ABCommC <- intersect(CMatA,CMatB)
  
  FinalC <- c(CMatA,CMatB)
  FinalR <- c(RMatA,RMatB)

  if(is.null(AMat) & !is.null(BMat))
  {
      return(BMat)
  }else if(is.null(BMat) & !is.null(AMat))
  {
      return(AMat)
  }

  if(length(ABCommC) == 0 & length(ABCommC) == 0)
  {
    print("merge two matrix!")
    ABMat <- matrix(0,nrow=dim(BMat)[1],ncol=dim(BMat)[2])
    BAMat <- matrix(0,nrow=dim(AMat)[1],ncol=dim(AMat)[2])
    
    FinalMatU <- cbind(AMat,ABMat)
    FinalMatD <- cbind(BAMat,BMat)
    FinalMat <- rbind(FinalMatU,FinalMatD)
    colnames(FinalMat) <- FinalC
    rownames(FinalMat) <- FinalR
    
    return(FinalMat)
  }
  else
  {
    print("merge two matrix fail!")
    return(FALSE)
  }
  
}


readCDNARawMat <- function(prefix=NULL,data.dir,ENSG = "F")
{
  #dir
  if (!dir.exists(data.dir))
  {
    stop(paste("Directory ",run," does not exist",sep=""))
  }
  if (!grepl("\\/$",data.dir))
  {
    data.dir <- paste(data.dir, "/", sep = "")
  }
  #dir
  
  #files path
  barcode.loc <- paste0(data.dir, "barcodes.tsv")
  gene.loc <- NULL
  gene.loc.1 <- paste0(data.dir,"genes.tsv")
  gene.loc.2 <- paste0(data.dir, "features.tsv")
  matrix.loc <- paste0(data.dir, "matrix.mtx")
  
  if (!file.exists(barcode.loc)) {
    stop("barcodes.tsv file is not existed!")
  }
  if (file.exists(gene.loc.1)) {
    print("find genes.tsv")
    gene.loc <- gene.loc.1
  }else if(file.exists(gene.loc.2)){
    print("find features.tsv")
    gene.loc <- gene.loc.2
  }else
  {
    stop("can't find genes.tsv or features.tsv")
  }
  
  if (!file.exists(matrix.loc)) {
    stop("matrix.mtx is not existed!")
  }
  
  data <- readMM(file = matrix.loc)
  cell.names <- readLines(barcode.loc)
  cell.names <- paste(prefix,cell.names,sep="")
  gene.names <- readLines(gene.loc)
 
  if(ENSG == "T")
  {
     GeneName <- make.unique(names = as.character(unlist(lapply(gene.names,function(x)unlist(strsplit(x,"\t"))[1]))))
  }else
  {
     GeneName <- make.unique(names = as.character(unlist(lapply(gene.names,function(x)unlist(strsplit(x,"\t"))[2]))))
  }

  rownames(data) <- GeneName
  colnames(data) <- cell.names
  
  return(data)
}

filterCellsNew <- function (x,cutoff=0.2,plot=TRUE)
{
  #filter Dead Cells
  mito.genes <- grep(pattern = "^MT-", x = rownames(x = x), value = TRUE)
  if(length(mito.genes) == 0){
    mito.genes <- grep(pattern = "^mt-", x = rownames(x = x), value = TRUE)
  }
  nUMI <- Matrix::colSums(x)
  percent.mito <- Matrix::colSums(x[mito.genes, ])/nUMI
  deadcells <- which(percent.mito>cutoff)
  filter.cells <- x[,-deadcells]
  print(dim(filter.cells))

  #filter Gene Exp is too big or too low
  x=filter.cells
  num.gene <- colSums(x > 0)
  filterGene <- names(x=num.gene[which(x = (num.gene> 200 & num.gene < 2500))])
  filCellGeneMat <- x[,filterGene]
  print(dim(filCellGeneMat))

  #filter Low Complexity Cells
  x=filCellGeneMat
  is.expr = 0
  nUMI <- Matrix::colSums(x)
  nGene <- colSums(x > is.expr)
  lm.x <- lm(nGene~nUMI)
  lm.x.res <- scale(lm.x$residuals)
  lowComCells <- which(lm.x.res[,1]< -3)
  filCellGeneLow <- x[,-lowComCells]
  print(dim(filCellGeneLow))
  
 
  return(filCellGeneLow)
}


CITESEQ <- function(RawMat,filBarList,CSeqNum,CSflag)
{
    GeneList <- rownames(RawMat)
    GeneNum <- dim(RawMat)[1]
    CITEBarList <- GeneList[(GeneNum-CSeqNum + 1):GeneNum]

    GCFMat <- RawMat[1:(GeneNum - CSeqNum),filBarList]

    CSBarExpTable <- cbind(filBarList,t(as.matrix(RawMat[CITEBarList,filBarList])))
    head <- c("cell",CITEBarList)
    colnames(CSBarExpTable) <- head
   
    return(CSBarExpTable) 
}

mergeMat <- function(AMat,BMat)
{
  AMat <- BMat
  return(AMat)
}

getFinalGCMat <- function(sample,CDNAListFile,flag,sbarFile,ENSG,OutBarFlag)
{
  #get all raw GC Mat
  RawGCMat <- NULL
  RawGCMatENSG <- NULL
  FilGCMat <- NULL
  AllSCITEBarExpMat <- NULL
  AllSCITEBarExpMatHead <- c()
  BarUMI <- c()

  readCDNAListResult <- readCDNAList(CDNAListFile)
  AllTypeList <- readCDNAListResult[[1]]
  CDNAdirList <- readCDNAListResult[[2]]
  CSflagList <- readCDNAListResult[[3]]

  for(i in seq(1,length(AllTypeList),1))
  {
    ThisType <- AllTypeList[i]
    CDNAdatadir = as.character(CDNAdirList[i])
    if(CSflagList[i] != "F")
    {
        CSeqNum = as.numeric(CSflagList[i])
    }else
    {
        CSeqNum = 0
    }

    stri <- paste(ThisType,CDNAdatadir,"CITESEQ:",CSeqNum,sep=",")
    print(stri)

    #row GC Mat
    CDNAFMat = readCDNARawMat(prefix=ThisType,CDNAdatadir,ENSG = ENSG)
    print("read matrix end!")
    
    #ENSG Matrix
    if(ENSG == "T")
    {
        CDNAFMatENSG <- readCDNARawMat(prefix=i,CDNAdatadir,ENSG = "T")
        RawGCMatENSG <- cbind(RawGCMatENSG,CDNAFMatENSG)
    }

    #CITESEQ
    if(CSeqNum != 0) #CITESEQ,cut mat,and get CITESEQBarMat
    {
        print("CSeqNum != 0")
        print(CSeqNum)
        GeneNum <- dim(CDNAFMat)[1]
        GeneName <- rownames(CDNAFMat)

        #cut raw mat and get citeseq normal result
        if(CSeqNum == 1)
        {
            ThisCSExpMat <- as.matrix(CDNAFMat[(GeneNum - CSeqNum + 1):GeneNum,])
        }else
        {
            ThisCSExpMat <- as.matrix(t(CDNAFMat[(GeneNum - CSeqNum + 1):GeneNum,]))
        }
        print(dim(ThisCSExpMat))
        AllSCITEBarExpMatHead <- c(AllSCITEBarExpMatHead,GeneName[(GeneNum - CSeqNum + 1):GeneNum])
        AllSCITEBarExpMatHead <- paste(AllSCITEBarExpMatHead,i,sep="-")

        CDNAFMat <- CDNAFMat[1:(GeneNum - CSeqNum),]
        BarUMI <- c(BarUMI,colSums(CDNAFMat))
    }

    #Filter
    if(flag == 1)
    {
        ThisFMat = NULL
    }else
    {
        ThisFMat = filterCellsNew(CDNAFMat,cutoff=0.2,plot=TRUE)
    }

    #merge before Matrix
    RawGCMat = cbind(RawGCMat,CDNAFMat)
    FilGCMat = cbind(FilGCMat,ThisFMat)

    #CITEBarExpMat merge with before sample
    if(CSeqNum != 0)
    {
        AllSCITEBarExpMat <- mergeTwoMat(AllSCITEBarExpMat,ThisCSExpMat)
    }

  }#deal every type sample end.

  #get need gc mat
  if(flag == 1) #get all raw matrix
  {
    stri <- paste(flag," - get all raw matrix!",sep="")
    GCMat <- RawGCMat
  }else if(flag == 2 | flag == 4) #get normal filter matrix
  {
    stri <- paste(flag," - get normal filter matrix!",sep="")
    GCMat <- FilGCMat
  }else if(flag == 5)
  {
    stri <- paste(flag," - only use special barcode!",sep="")
    tolbar <- colnames(RawGCMat)
    spebar <- readLines(sbarFile)
    combar <- intersect(spebar,tolbar)
    
    GCMat <- RawGCMat[,combar]
    stri <- paste(stri,",now barcode:",length(combar),sep="")
    print(stri)
  }
  else if(flag == 3) #get normal filter + special bar matrix
  {
    stri <- paste(flag," - get normal filter + special bar matrix!",sep="")
    spebar <- readLines(sbarFile)
    filbar <- colnames(FilGCMat)
    nabar <- setdiff(spebar,filbar) #need add barcode
    tolbar <- colnames(RawGCMat)
    addbar <- intersect(tolbar,nabar) #final add barcode
    Nobar <- setdiff(nabar,addbar)  #CDNA could't seq some barcode come from TCR or BCR!!!
    stri <- paste(stri,"add special barcode:",length(addbar),sep="")
    GCMat <- cbind(FilGCMat,RawGCMat[,addbar])
    stri <- paste(stri,",CDNA no have barcode:",length(Nobar),sep="")

    if(length(Nobar) > 0)
    {
      GCMatRawHead <- colnames(GCMat)
      GCMatNowHead <- c(GCMatRawHead,Nobar)
      GCMat <- cbind(GCMat,matrix(0,nrow = dim(GCMat)[1],ncol = length(Nobar)))
      colnames(GCMat) <- GCMatNowHead
    }
    
  }
  print(stri)
 
  #get ENSG Matrix 
  if(ENSG == "T")
  {
     print("ENSG Matrix!")
     finalBar <- colnames(GCMat)
     GCMat <- RawGCMatENSG[,finalBar]
  }

  #output
  matfile <- paste(sample,".GCMat.RData",sep = "")
  print(dim(GCMat))
  save(GCMat,file=matfile)
  stri <- paste("output GCMat:",matfile,sep="")
  print(stri)

  if(OutBarFlag == "T")
  {
     OutBarFile <- paste(sample,".Barcode.txt",sep = "")
     AllBar <- colnames(GCMat)
     writeLines(AllBar,con=OutBarFile,sep="\n")
     stri <- paste("output Barcode:",OutBarFile,sep="")
     print(stri)
  }

  if(CSeqNum != 0)
  {
     print("save CITESEQ Barcode Exp GCMat + Table RData:") #col is barcode,row is CITESEQ barcode
     CITEExpTableRData <- paste(sample,".CITEBarcodeExpTable.RData",sep = "")
     colnames(AllSCITEBarExpMat) <- AllSCITEBarExpMatHead
     save(AllSCITEBarExpMat,file=CITEExpTableRData)

  }

  print("finish!")
}

#main
setwd(workDir)
getFinalGCMat(sample,CDNAlist,flag,SpeBarFile,ENSGflag,OutBarFlag)

  
